#include <jni.h>
#include <string>
#include "aes.h"

using namespace std;

const char *AES_ALGORITHM = "AES/CBC/PKCS5Padding";

jbyteArray encryption(JNIEnv *env, jbyteArray srcData, string secretKey,string secretIv) {
    jstring type = (*env).NewStringUTF("AES");
    jstring cipher_mode = (*env).NewStringUTF(AES_ALGORITHM);
    jbyteArray pwd = (*env).NewByteArray(16);
    char *master_key = (char *) secretKey.c_str();
    (*env).SetByteArrayRegion(pwd, 0, 16, reinterpret_cast<jbyte *>(master_key));
    // 向量
    jbyteArray iv = (*env).NewByteArray(16);
    char *iv_key = (char *) secretIv.c_str();
    (*env).SetByteArrayRegion(iv, 0, 16, reinterpret_cast<jbyte *>(iv_key));

    jclass secretKeySpecClass = (*env).FindClass("javax/crypto/spec/SecretKeySpec");
    jmethodID secretKeySpecMethodId = (*env).GetMethodID(secretKeySpecClass, "<init>",
                                                         "([BLjava/lang/String;)V");
    jobject secretKeySpecObj = (*env).NewObject(secretKeySpecClass, secretKeySpecMethodId, pwd,
                                                type);
    // 创建IvParameterSpec
    jclass ivParameterSpecClass = (*env).FindClass("javax/crypto/spec/IvParameterSpec");
    jmethodID ivParameterSpecMethodId = (*env).GetMethodID(ivParameterSpecClass, "<init>",
                                                           "([B)V");
    jobject ivParameterSpecObj = (*env).NewObject(ivParameterSpecClass, ivParameterSpecMethodId,
                                                  iv);
    jclass cipherClass = (*env).FindClass("javax/crypto/Cipher");
    jmethodID cipherInitMethodId = (*env).GetMethodID(cipherClass, "init",
                                                      "(ILjava/security/Key;Ljava/security/spec/AlgorithmParameterSpec;)V");
    jmethodID cipherInstanceMethodId = (*env).GetStaticMethodID(cipherClass, "getInstance",
                                                                "(Ljava/lang/String;)Ljavax/crypto/Cipher;");
    jobject cipherObj = (*env).CallStaticObjectMethod(cipherClass, cipherInstanceMethodId,
                                                      cipher_mode);
    jfieldID decryptModeFieldId = (*env).GetStaticFieldID(cipherClass, "ENCRYPT_MODE", "I");
    jint mode = (*env).GetStaticIntField(cipherClass, decryptModeFieldId);
    (*env).CallVoidMethod(cipherObj, cipherInitMethodId, mode, secretKeySpecObj,
                          ivParameterSpecObj);
    jmethodID doFinalMethodId = (*env).GetMethodID(cipherClass, "doFinal", "([B)[B");
    jbyteArray text = (jbyteArray) (*env).CallObjectMethod(cipherObj, doFinalMethodId, srcData);
    return text;
}

jbyteArray decryption(JNIEnv *env, jbyteArray srcData, string secretKey,string secretIv) {
    jstring type = (*env).NewStringUTF("AES");
    jstring cipher_mode = (*env).NewStringUTF(AES_ALGORITHM);
    jbyteArray pwd = (*env).NewByteArray(16);
    char *master_key = (char *) secretKey.c_str();
    (*env).SetByteArrayRegion(pwd, 0, 16, reinterpret_cast<jbyte *>(master_key));
    // 向量
    jbyteArray iv = (*env).NewByteArray(16);
    char *iv_key = (char *) secretIv.c_str();
    (*env).SetByteArrayRegion(iv, 0, 16, reinterpret_cast<jbyte *>(iv_key));
    jclass secretKeySpecClass = (*env).FindClass("javax/crypto/spec/SecretKeySpec");
    jmethodID secretKeySpecMethodId = (*env).GetMethodID(secretKeySpecClass, "<init>",
                                                         "([BLjava/lang/String;)V");
    jobject secretKeySpecObj = (*env).NewObject(secretKeySpecClass, secretKeySpecMethodId, pwd,
                                                type);
    // 创建IvParameterSpec
    jclass ivParameterSpecClass = (*env).FindClass("javax/crypto/spec/IvParameterSpec");
    jmethodID ivParameterSpecMethodId = (*env).GetMethodID(ivParameterSpecClass, "<init>",
                                                           "([B)V");
    jobject ivParameterSpecObj = (*env).NewObject(ivParameterSpecClass, ivParameterSpecMethodId,
                                                  iv);
    jclass cipherClass = (*env).FindClass("javax/crypto/Cipher");
    jmethodID cipherInitMethodId = (*env).GetMethodID(cipherClass, "init",
                                                      "(ILjava/security/Key;Ljava/security/spec/AlgorithmParameterSpec;)V");
    jmethodID cipherInstanceMethodId = (*env).GetStaticMethodID(cipherClass, "getInstance",
                                                                "(Ljava/lang/String;)Ljavax/crypto/Cipher;");
    jobject cipherObj = (*env).CallStaticObjectMethod(cipherClass, cipherInstanceMethodId,
                                                      cipher_mode);
    jfieldID decryptModeFieldId = (*env).GetStaticFieldID(cipherClass, "DECRYPT_MODE", "I");
    jint mode = (*env).GetStaticIntField(cipherClass, decryptModeFieldId);
    (*env).CallVoidMethod(cipherObj, cipherInitMethodId, mode, secretKeySpecObj,
                          ivParameterSpecObj);
    jmethodID doFinalMethodId = (*env).GetMethodID(cipherClass, "doFinal", "([B)[B");
    jbyteArray text = (jbyteArray) (*env).CallObjectMethod(cipherObj, doFinalMethodId, srcData);
    return text;
}


